////////////////////////////////////////////////////////////////////////
// a3.cpp
//
// Third assignment for Advanced Multimedia Computing
//
// T. Kostelijk & F. Huizinga
//     0418889      0418862
//
////////////////////////////////////////////////////////////////////////
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cv.h>
#include <highgui.h>
#include <vector>
#include <math.h>
#include <time.h>

CvMemStorage *storage;
CvHaarClassifierCascade *cascade;
CvFont font;
// initialise colors
CvScalar colorOrange = cvScalar(0,188,255);
CvScalar colorBlue = cvScalar(256,0,0);
CvScalar colorGreen = cvScalar(0,256,0);
CvScalar colorRed = cvScalar(0,0,256);
CvScalar colorBlack = cvScalar(0,0,0);
CvScalar colorWhite = cvScalar(256,256,256);
int counter       = 0;

/* Show and save image to harddisk */
inline void showAndSave(IplImage *img) {
	char buf[256];
	sprintf(buf, "%8.8d.png", counter);
	if (!cvSaveImage(buf, img))
		printf("Could not save image\n");
	else
		counter++;
	cvShowImage("mainWin", img);
}

struct HaarObject{
	unsigned x, y, w, h, startFrame, nrPreds, faceId;

	HaarObject() {
	}
	HaarObject(unsigned x, unsigned y, unsigned w, unsigned h,unsigned startFrame, unsigned nrPreds=0, unsigned faceId=0) {
		this->x = x;
		this->y = y;
		this->w = w;
		this->h = h;
		this->startFrame = startFrame;
		this->nrPreds = 0;
		this->faceId = faceId;
	}

	void printme(){
		printf("x:%d \ty:%d \tw:%d \th:%d \tstartFrame:%d \tfaceId:%d nrPreds:%d\n", this->x, this->y, this->w, this->h, this->startFrame, this->faceId, this->nrPreds);
	}
};

// get difference in location of two haar object by calculating the euclidean distance
float getLocDiff(HaarObject haarObj1, HaarObject haarObj2) {
	int deltaX = abs(haarObj1.x-haarObj2.x);
	int deltaY = abs(haarObj1.y-haarObj2.y);
	return sqrt((deltaX*deltaX)+(deltaY*deltaY));
}
// get size difference of two haar objects
float getSizeDiff(HaarObject haarObj1, HaarObject haarObj2) {
	return abs(sqrt((haarObj1.w*haarObj1.w)+(haarObj1.h*haarObj1.h)) - sqrt((haarObj2.w*haarObj2.w)+(haarObj2.h*haarObj2.h)));
}

// prints the vector of haarobjectvectors (debug purpose)
void printVectS( std::vector< std::vector<HaarObject> > &haarObjVectS ) {
	printf("printing haarObjVectS ..\n");
	if(haarObjVectS.size() == 0) {
		printf("object empty!\n");
	}
	for(unsigned i=0;i<haarObjVectS.size();i++){
		printf("i=%d -----------------------------------------------------------------\n",i);
		for(unsigned j=0;j<haarObjVectS[i].size();j++){
			haarObjVectS[i][j].printme();
		}
	}
}

void detectAndDraw(IplImage *img, float scale, std::vector< std::vector<HaarObject> > &haarObjVectS) {
	// create empty image
	IplImage *gray = cvCreateImage(cvSize(img->width, img->height), 8, 1 );
	// scale the empty image to smaller format
	IplImage *smallImg = cvCreateImage(cvSize(img->width/scale, img->height/scale), 8, 1);
	// convert to gray
	cvCvtColor(img, gray, CV_BGR2GRAY );
	// resize image with scale factor
	cvResize(gray, smallImg, CV_INTER_LINEAR);
	// maximise contrast 
	cvEqualizeHist(smallImg, smallImg);
	// clear the memory
	cvClearMemStorage(storage);
	// initalise haar ector
	std::vector<HaarObject> haarObjVect;

	//detect the objects
	CvSeq *objects = cvHaarDetectObjects(smallImg, cascade, storage, 1.1, 1, CV_HAAR_DO_CANNY_PRUNING, cvSize(1,1));

	// loop through objects
	for(int i=0; i < (objects ? objects->total : 0); i++) {
		// get object
		CvRect *r = (CvRect*) cvGetSeqElem(objects, i);
		//cvRectangle(img, cvPoint(r->x*scale, r->y*scale), cvPoint((r->x+r->width)*scale, (r->y+r->height)*scale), colorBlack);

		HaarObject haarObj = HaarObject(r->x, r->y, r->width, r->height, haarObjVectS.size());
		// store haar object
		haarObjVect.push_back(haarObj);
	}
	haarObjVectS.push_back(haarObjVect);
	printVectS(haarObjVectS);

	//showAndSave(img);

	// write image

	// clear the memory
	cvReleaseImage( &gray );
	cvReleaseImage( &smallImg );
}

// predicts new face from two locations by 0th order derivative
HaarObject getPrediction(HaarObject &A, HaarObject &B) {
	HaarObject R;
	// add delta distance 
	R.x = B.x + (B.x - A.x);
	R.y = B.y + (B.y - A.y);
	R.w = B.w;
	R.h = B.h;
	R.faceId = A.faceId;
	R.nrPreds = A.nrPreds + 1;
	return R;
}

// draws a rectangle of a haarobject with its face id
void drawRectText(IplImage *img, HaarObject &hObj, CvScalar c) {
	char faceIdChar[255];
	// todo scale
	// draw face
	unsigned X = hObj.x;
	unsigned Y = hObj.y;
	unsigned W = X+hObj.w;
	unsigned H = Y+hObj.h;
	cvRectangle(img, cvPoint(X,Y), cvPoint(W,H), c);
	sprintf(faceIdChar,"%d", hObj.faceId);
	//sprintf(faceIdChar,"I:%d, L:%d,S:%d", haarObjVectS[lastFrame][y].faceId, locDiff, sizeDiff);
	//printf("faceIdChar : %s\n", faceIdChar);
	cvRectangle(img, cvPoint(X,Y-10), cvPoint(X+30,Y), colorBlack, CV_FILLED);
	cvPutText(img, faceIdChar, cvPoint(X,Y), &font, c);
}

// detect faces that are stable, i.e. faces that continue are in neighbor of the predictions
unsigned detectStableFace(IplImage *img, std::vector< std::vector<HaarObject> > &haarObjVectS, std::vector< std::vector<HaarObject> > &predVectS, float scale, unsigned maxFaceId) {
	unsigned locDiff,sizeDiff,faceId;
	//unsigned maxLocDiff = 6;
	unsigned lastFrame = haarObjVectS.size()-1;
	float X,Y,W,H;
	bool bBreak=false,bPredFound;
	HaarObject pred;
	std::vector<HaarObject> predVect,v;
	unsigned a,b,c,j,k,y,y2;
	int i;
	unsigned nrPreds;
	unsigned faceCount;
	int minFaceLast=4;

	// tweakable parameters
	const unsigned hist=1,maxLocDiff = 15, maxRect=60, minNrPreds=4, maxSizeDiff = 7;

	// copy all objects in first frame to prediction objects
	printf("lastFrame = %d\n",lastFrame);

	// first frame
	if (lastFrame==0) {
		// We are in the first frame so the predictions are simply the found haarobjects
		predVectS.push_back( haarObjVectS[lastFrame] );
		// Loop through the prediction vector and hand out face id's
		for (b=0;b<predVectS[lastFrame].size();b++){
			predVectS[lastFrame][b].faceId = b+1;
		}	
		// keep track on the handed out face id's by storing it in maxFaceId
		maxFaceId = predVectS[lastFrame].size();
	}

	// draw predictions (orange rectangles)
	//cvWaitKey(0);
	/*
	a = predVectS.size()-1;
	for (b=0;b<predVectS[a].size();b++){
		drawRectText(img, predVectS[a][b], colorOrange);
		showAndSave(img);
	}	
	*/
	//cvWaitKey(0);

	// make use of last prediction, thus previous frame
	i = predVectS.size()-1;

	// loop through objects of last frame, (ignoring first frame and to big faces)
	for(y=0; y<haarObjVectS[lastFrame].size() && lastFrame>0 && (haarObjVectS[lastFrame][y].w<maxRect && haarObjVectS[lastFrame][y].h<maxRect); y++){
		bPredFound = false;
		// loop through predicted objects
		for(j=0;j<predVectS[i].size() && !bPredFound ;j++){
			locDiff = getLocDiff(haarObjVectS[lastFrame][y], predVectS[i][j]);
			// draw rectange of the face we are evaluating
			//drawRectText(img, haarObjVectS[lastFrame][y], colorRed);
			// draw rectange of the prediction we are comparing to 
			//drawRectText(img, predVectS[i][j], colorRed);
			//showAndSave(img);
			printf("[start compare] y=%d j=%d\n", y, j);
			haarObjVectS[lastFrame][y].printme();
			predVectS[i][j].printme();
			printf("locDiff = %d should be < %d\n",locDiff,maxLocDiff);
			//drawRectText(img, haarObjVectS[lastFrame][y], colorBlack);
			//drawRectText(img, predVectS[i][j], colorOrange);
			// if the location difference is smaller then the maximum difference (manualy tweaked param)
			if(locDiff<maxLocDiff) {
				//cvWaitKey(1500);
				bPredFound = true;
				printf(">>FOUND nearby face with faceId = %d\n",predVectS[i][j].faceId); 
				// if prediction has an existing face id
				if (predVectS[i][j].faceId != 0) {
					printf("overloading faceId\n");
					faceId = predVectS[i][j].faceId;
					nrPreds = predVectS[i][j].nrPreds;
				} else {
					// create new face id
					maxFaceId++;
					nrPreds = 0;
					faceId = maxFaceId;
					printf("NOT overloading faceId, creating new face id : %d\n", faceId);
				}
				// setting face id and number of times this face is predicted
				haarObjVectS[lastFrame][y].faceId = faceId;
				haarObjVectS[lastFrame][y].nrPreds = nrPreds;
				pred = getPrediction(predVectS[i][j], haarObjVectS[lastFrame][y]);
				printf("PREDICTION = \n");
				pred.printme(); 
			}
		}
		if(bPredFound) {
			predVectS.push_back(predVect);
		} else {
			printf("no pred found\n");
			pred = haarObjVectS[lastFrame][y];
			pred.faceId = ++maxFaceId;
		}
		predVect.push_back( pred );
	}
	predVectS.push_back(predVect);


	faceCount = 0;
	// count nr Faces

	unsigned faceCounts[maxFaceId+1];
	// initialise all on 0
	for (i=0; i<maxFaceId+1; i++) { faceCounts[i] = 0; }

	// store the number of predictions per faceid (which can be used as a measure of stability)
	for (i=0;i<predVectS.size();i++){
		for (j=0;j<predVectS[i].size();j++){
			faceCounts[predVectS[i][j].faceId] = predVectS[i][j].nrPreds;
		}	
	}
	i--;
	for (j=0;j<predVectS[i].size();j++){
		// draw rectangles on faces which last minimal 2 frames
		if(predVectS[i][j].nrPreds>=minFaceLast){
			drawRectText(img, predVectS[i][j], colorRed);
		}
	}	
	
	// if a face is predicted more then 2 times its stable 
	for (i=0; i<maxFaceId+1; i++) { 
		printf("array faceCounts[%d] = %d",i,faceCounts[i]); 
		if(faceCounts[i] > minFaceLast) {
			faceCount++;
		}
	}
	char faceCountChar[255];
	sprintf(faceCountChar,"%d", faceCount);
	// print facecount
	cvRectangle(img, cvPoint(10,10), cvPoint(30,20), colorBlack, CV_FILLED);
	cvPutText(img, faceCountChar, cvPoint(10,20), &font, colorRed);
	//showAndSave(img);
	return maxFaceId;
}

int main(int argc, char **argv) {
	IplImage *img;
	char path[] = "data/movie2/";
	//char path[] = "data/movie2/";
	char fileName[256];
	char pathFileName[256];
	std::vector< std::vector<HaarObject> > haarObjVectS;
	std::vector< std::vector<HaarObject> > predVectS;
	float scale=1;
	unsigned maxFaceId=0;

	// initialize font
	cvInitFont(&font, CV_FONT_HERSHEY_PLAIN, 1.0, 1.0, 0, 1, 8);
	// initialize storage
	storage = cvCreateMemStorage(0);
	cascade = (CvHaarClassifierCascade*) cvLoad("/usr/share/opencv/haarcascades/haarcascade_frontalface_alt2.xml", 0, 0, 0);

	// start window
	cvNamedWindow("mainWin", CV_WINDOW_AUTOSIZE); 

	for (int i=1;i<=285;i++){
	//for (int i=1;i<=319;i++){
		printf("frame %d\n",i);
		printf("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n");
		sprintf(fileName, "%8.8d.png", i); 
		sprintf(pathFileName, "%s%s", path, fileName);
		img = cvLoadImage(pathFileName);
		if(!img){
			printf("Could not load image file: %s\n",fileName);
			exit(0);
		}

		printf("detecting faces with haar..\n");
		// detect the faces 
		detectAndDraw(img, 1, haarObjVectS);

		printf("detecting stable faces..\n");
		// detect the faces over a few frames
		maxFaceId = detectStableFace(img, haarObjVectS, predVectS, scale, maxFaceId);
		printf(".............stable DONE\n");

		showAndSave(img);
		//if(!cvSaveImage(fileName, img)) printf("Could not save: %s\n", fileName);
		//printVectS( haarObjVectS );
		//showAndSave(img);
		//cvWaitKey(0);
	}

	//cvWaitKey(0);
	// clear the memory
	cvReleaseImage(&img);
	return 0;
}

